package io.kwdb.bench;

import io.kwdb.sdk.*;
import io.kwdb.sdk.internal.Metrics;
import java.io.*;
import java.lang.management.GarbageCollectorMXBean;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.sql.*;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Logger;

/**
 * Performance benchmark runner for KWDB SDK. Supports gRPC, ILP, and JDBC transports.
 *
 * <p>Usage: --transport={grpc|ilp|jdbc} --threads=4 --batch=500 --lingerMs=10 --duration=60s
 * --warmup=10s --measurement=sdk_bench --fields=5 --tags=host,region --rateLimit=0 --timeoutMs=3000
 */
public class BenchRunner {

  private static final Logger logger = Logger.getLogger(BenchRunner.class.getName());
  private static final Random random = new Random(42); // Fixed seed for reproducibility

  // Command line arguments
  private String transport = "grpc";
  private int threads = 4;
  private int batchSize = 500;
  private int lingerMs = 10;
  private int durationSeconds = 60;
  private int warmupSeconds = 10;
  private String measurement = "sdk_bench";
  private int numFields = 5;
  private String tagNames = "host,region";
  private int rateLimit = 0; // 0 = unlimited
  private int timeoutMs = 3000;
  private String csvOutput = null;
  private boolean verify = true;

  // Test configuration
  private String dbHost = "127.0.0.1";
  private int grpcPort = 9876;
  private int httpPort = 9091;
  private String jdbcUrl = "jdbc:postgresql://127.0.0.1:5432/kwdb";
  private String jdbcUser = "kwdb";
  private String jdbcPassword = "kwdb";

  // Runtime state
  private final AtomicLong totalPoints = new AtomicLong(0);
  private final AtomicLong errorCount = new AtomicLong(0);
  private final AtomicInteger retryCount = new AtomicInteger(0);
  private final List<Long> latencies = Collections.synchronizedList(new ArrayList<>());
  private volatile boolean running = false;
  private long startTime;
  private long endTime;

  public static void main(String[] args) throws Exception {
    BenchRunner runner = new BenchRunner();
    runner.parseArgs(args);
    runner.run();
  }

  private void parseArgs(String[] args) {
    for (String arg : args) {
      if (arg.startsWith("--transport=")) {
        transport = arg.substring("--transport=".length());
      } else if (arg.startsWith("--threads=")) {
        threads = Integer.parseInt(arg.substring("--threads=".length()));
      } else if (arg.startsWith("--batch=")) {
        batchSize = Integer.parseInt(arg.substring("--batch=".length()));
      } else if (arg.startsWith("--lingerMs=")) {
        lingerMs = Integer.parseInt(arg.substring("--lingerMs=".length()));
      } else if (arg.startsWith("--duration=")) {
        String duration = arg.substring("--duration=".length());
        durationSeconds = parseDuration(duration);
      } else if (arg.startsWith("--warmup=")) {
        String warmup = arg.substring("--warmup=".length());
        warmupSeconds = parseDuration(warmup);
      } else if (arg.startsWith("--measurement=")) {
        measurement = arg.substring("--measurement=".length());
      } else if (arg.startsWith("--fields=")) {
        numFields = Integer.parseInt(arg.substring("--fields=".length()));
      } else if (arg.startsWith("--tags=")) {
        tagNames = arg.substring("--tags=".length());
      } else if (arg.startsWith("--rateLimit=")) {
        rateLimit = Integer.parseInt(arg.substring("--rateLimit=".length()));
      } else if (arg.startsWith("--timeoutMs=")) {
        timeoutMs = Integer.parseInt(arg.substring("--timeoutMs=".length()));
      } else if (arg.startsWith("--csvOutput=")) {
        csvOutput = arg.substring("--csvOutput=".length());
      } else if (arg.startsWith("--verify=")) {
        verify = Boolean.parseBoolean(arg.substring("--verify=".length()));
      }
    }
  }

  private int parseDuration(String duration) {
    if (duration.endsWith("s")) {
      return Integer.parseInt(duration.substring(0, duration.length() - 1));
    } else if (duration.endsWith("m")) {
      return Integer.parseInt(duration.substring(0, duration.length() - 1)) * 60;
    } else {
      return Integer.parseInt(duration);
    }
  }

  private void run() throws Exception {
    printConfiguration();

    // Reset metrics
    Metrics.getInstance().reset();

    // Warmup phase
    if (warmupSeconds > 0) {
      System.out.println("\n=== WARMUP PHASE (" + warmupSeconds + "s) ===");
      runBenchmark(warmupSeconds * 1000L, true);

      // Reset counters after warmup
      totalPoints.set(0);
      errorCount.set(0);
      retryCount.set(0);
      latencies.clear();
      Metrics.getInstance().reset();
    }

    // Measurement phase
    System.out.println("\n=== MEASUREMENT PHASE (" + durationSeconds + "s) ===");
    startTime = System.currentTimeMillis();
    runBenchmark(durationSeconds * 1000L, false);
    endTime = System.currentTimeMillis();

    // Print results
    printResults();

    // Verify if requested
    if (verify && !transport.equals("jdbc")) {
      verifyData();
    }

    // Write CSV if requested
    if (csvOutput != null) {
      writeCsv();
    }
  }

  private void printConfiguration() {
    System.out.println("\n" + "=".repeat(70));
    System.out.println(" KWDB SDK BENCHMARK - M4");
    System.out.println("=".repeat(70));
    System.out.println("Configuration:");
    System.out.println("  Transport: " + transport);
    System.out.println("  Threads: " + threads);
    System.out.println("  Batch Size: " + batchSize);
    System.out.println("  Linger MS: " + lingerMs);
    System.out.println("  Duration: " + durationSeconds + "s");
    System.out.println("  Warmup: " + warmupSeconds + "s");
    System.out.println("  Measurement: " + measurement);
    System.out.println("  Fields: " + numFields);
    System.out.println("  Tags: " + tagNames);
    System.out.println("  Rate Limit: " + (rateLimit == 0 ? "unlimited" : rateLimit + " points/s"));
    System.out.println("  Timeout: " + timeoutMs + "ms");
    System.out.println("=".repeat(70));
  }

  private void runBenchmark(long durationMs, boolean isWarmup) throws Exception {
    running = true;
    ExecutorService executor = Executors.newFixedThreadPool(threads);
    List<Future<?>> futures = new ArrayList<>();

    // Rate limiter (if enabled)
    RateLimiter rateLimiter = rateLimit > 0 ? new RateLimiter(rateLimit) : null;

    // Start worker threads
    for (int t = 0; t < threads; t++) {
      final int threadId = t;
      futures.add(
          executor.submit(
              () -> {
                try {
                  if (transport.equals("jdbc")) {
                    runJdbcWorker(threadId, durationMs, isWarmup, rateLimiter);
                  } else {
                    runSdkWorker(threadId, durationMs, isWarmup, rateLimiter);
                  }
                } catch (Exception e) {
                  logger.severe("Worker " + threadId + " failed: " + e.getMessage());
                  e.printStackTrace();
                }
              }));
    }

    // Monitor progress
    long startTime = System.currentTimeMillis();
    while (System.currentTimeMillis() - startTime < durationMs) {
      Thread.sleep(1000);
      if (!isWarmup) {
        long pointsSoFar = totalPoints.get();
        long elapsed = (System.currentTimeMillis() - startTime) / 1000;
        if (elapsed > 0) {
          System.out.printf(
              "Progress: %d points, %d points/s, %d errors\r",
              pointsSoFar, pointsSoFar / elapsed, errorCount.get());
        }
      }
    }

    // Stop workers
    running = false;
    executor.shutdown();
    executor.awaitTermination(10, TimeUnit.SECONDS);

    // Wait for all futures
    for (Future<?> future : futures) {
      try {
        future.get(5, TimeUnit.SECONDS);
      } catch (TimeoutException e) {
        future.cancel(true);
      }
    }
  }

  private void runSdkWorker(
      int threadId, long durationMs, boolean isWarmup, RateLimiter rateLimiter) throws Exception {
    // Create client based on transport
    Config config;

    if (transport.equals("grpc")) {
      config =
          Config.builder()
              .host(dbHost)
              .port(grpcPort)
              .transport(Config.Transport.GRPC)
              .timeoutMillis(timeoutMs)
              .maxBatchSize(batchSize)
              .lingerMs(lingerMs)
              .enableRetry(true)
              .maxRetries(2)
              .build();
    } else if (transport.equals("ilp")) {
      config =
          Config.builder()
              .host(dbHost)
              .httpPort(httpPort)
              .transport(Config.Transport.ILP)
              .timeoutMillis(timeoutMs)
              .maxBatchSize(batchSize)
              .lingerMs(lingerMs)
              .enableRetry(true)
              .maxRetries(2)
              .build();
    } else {
      throw new IllegalArgumentException("Unknown transport: " + transport);
    }

    try (KwdbClient client = KwdbClient.create(config)) {
      long workerStart = System.currentTimeMillis();
      String[] tags = tagNames.split(",");

      while (running && (System.currentTimeMillis() - workerStart < durationMs)) {
        // Rate limiting
        if (rateLimiter != null) {
          rateLimiter.acquire(batchSize);
        }

        // Create batch
        List<Point> batch = new ArrayList<>();
        for (int i = 0; i < batchSize; i++) {
          Point.Builder builder = Point.builder(measurement);

          // Add tags
          for (String tag : tags) {
            if (tag.equals("host")) {
              builder.tag("host", "server-" + (threadId % 10));
            } else if (tag.equals("region")) {
              builder.tag("region", threadId % 2 == 0 ? "us-west" : "us-east");
            } else {
              builder.tag(tag, "value-" + random.nextInt(10));
            }
          }

          // Add fields
          for (int f = 0; f < numFields; f++) {
            String fieldName = "f" + f;
            switch (f % 4) {
              case 0:
                builder.field(fieldName, random.nextDouble() * 100);
                break;
              case 1:
                builder.field(fieldName, (long) (random.nextInt(1000)));
                break;
              case 2:
                builder.field(fieldName, random.nextBoolean());
                break;
              case 3:
                builder.field(fieldName, "str-" + random.nextInt(100));
                break;
            }
          }

          builder.timestamp(System.currentTimeMillis() * 1_000_000L + i);
          batch.add(builder.build());
        }

        // Write batch
        long writeStart = System.currentTimeMillis();
        Status status = client.batchWrite(batch);
        long writeLatency = System.currentTimeMillis() - writeStart;

        if (!isWarmup) {
          if (status.isOk()) {
            totalPoints.addAndGet(batchSize);
            latencies.add(writeLatency);
          } else {
            errorCount.addAndGet(batchSize);
          }
        }
      }
    }
  }

  private void runJdbcWorker(
      int threadId, long durationMs, boolean isWarmup, RateLimiter rateLimiter) throws Exception {
    // JDBC implementation
    try (Connection conn = DriverManager.getConnection(jdbcUrl, jdbcUser, jdbcPassword)) {
      conn.setAutoCommit(false);

      // Prepare statement for batch insert
      String[] tags = tagNames.split(",");
      StringBuilder sql = new StringBuilder("INSERT INTO " + measurement + " (");

      // Build column list
      for (String tag : tags) {
        sql.append(tag).append(", ");
      }
      for (int f = 0; f < numFields; f++) {
        sql.append("f").append(f).append(", ");
      }
      sql.append("ts) VALUES (");

      // Build placeholders
      for (int i = 0; i < tags.length + numFields; i++) {
        sql.append("?, ");
      }
      sql.append("?)");

      PreparedStatement pstmt = conn.prepareStatement(sql.toString());
      long workerStart = System.currentTimeMillis();

      while (running && (System.currentTimeMillis() - workerStart < durationMs)) {
        // Rate limiting
        if (rateLimiter != null) {
          rateLimiter.acquire(batchSize);
        }

        long writeStart = System.currentTimeMillis();

        // Create batch
        for (int i = 0; i < batchSize; i++) {
          int paramIndex = 1;

          // Set tags
          for (String tag : tags) {
            if (tag.equals("host")) {
              pstmt.setString(paramIndex++, "server-" + (threadId % 10));
            } else if (tag.equals("region")) {
              pstmt.setString(paramIndex++, threadId % 2 == 0 ? "us-west" : "us-east");
            } else {
              pstmt.setString(paramIndex++, "value-" + random.nextInt(10));
            }
          }

          // Set fields
          for (int f = 0; f < numFields; f++) {
            switch (f % 4) {
              case 0:
                pstmt.setDouble(paramIndex++, random.nextDouble() * 100);
                break;
              case 1:
                pstmt.setLong(paramIndex++, random.nextInt(1000));
                break;
              case 2:
                pstmt.setBoolean(paramIndex++, random.nextBoolean());
                break;
              case 3:
                pstmt.setString(paramIndex++, "str-" + random.nextInt(100));
                break;
            }
          }

          // Set timestamp
          pstmt.setLong(paramIndex++, System.currentTimeMillis() * 1_000_000L + i);
          pstmt.addBatch();
        }

        // Execute batch
        try {
          int[] results = pstmt.executeBatch();
          conn.commit();

          long writeLatency = System.currentTimeMillis() - writeStart;

          if (!isWarmup) {
            totalPoints.addAndGet(batchSize);
            latencies.add(writeLatency);
          }
        } catch (SQLException e) {
          conn.rollback();
          if (!isWarmup) {
            errorCount.addAndGet(batchSize);
          }
        }
      }
    }
  }

  private void printResults() {
    long duration = (endTime - startTime) / 1000;
    long totalWritten = totalPoints.get();
    long throughput = duration > 0 ? totalWritten / duration : 0;
    double errorRate = totalWritten > 0 ? (double) errorCount.get() / totalWritten * 100 : 0;

    // Calculate percentiles
    Collections.sort(latencies);
    long p50 = calculatePercentile(latencies, 50);
    long p95 = calculatePercentile(latencies, 95);
    long p99 = calculatePercentile(latencies, 99);

    // Get GC stats
    long gcCount = 0;
    for (GarbageCollectorMXBean gc : ManagementFactory.getGarbageCollectorMXBeans()) {
      gcCount += gc.getCollectionCount();
    }

    // Get memory stats
    MemoryMXBean memory = ManagementFactory.getMemoryMXBean();
    long heapMb = memory.getHeapMemoryUsage().getUsed() / (1024 * 1024);

    System.out.println("\n\n=== RESULTS ===");
    System.out.println("Transport: " + transport);
    System.out.println("Duration: " + duration + "s");
    System.out.println("Total Points: " + totalWritten);
    System.out.println("Throughput: " + throughput + " points/s");
    System.out.println("Latency P50: " + p50 + "ms");
    System.out.println("Latency P95: " + p95 + "ms");
    System.out.println("Latency P99: " + p99 + "ms");
    System.out.println("Error Rate: " + String.format("%.2f%%", errorRate));
    System.out.println("Retries: " + retryCount.get());
    System.out.println("GC Count: " + gcCount);
    System.out.println("Heap Used: " + heapMb + "MB");
    System.out.println("=".repeat(50));
  }

  private long calculatePercentile(List<Long> sorted, int percentile) {
    if (sorted.isEmpty()) return 0;
    int index = (int) Math.ceil(percentile / 100.0 * sorted.size()) - 1;
    return sorted.get(Math.max(0, Math.min(index, sorted.size() - 1)));
  }

  private void verifyData() {
    try {
      Config config =
          Config.builder()
              .host(dbHost)
              .port(grpcPort)
              .transport(Config.Transport.GRPC)
              .timeoutMillis(5000)
              .build();

      try (KwdbClient client = KwdbClient.create(config)) {
        String sql = "SELECT COUNT(*) as cnt FROM " + measurement;
        QueryResult result = client.query(sql);

        if (result.isSuccess() && result.getRowCount() > 0) {
          Long count = result.getLong(0, "cnt");
          System.out.println("\nVerification: Database contains " + count + " points");
          System.out.println("Expected: ~" + totalPoints.get() + " points");

          if (Math.abs(count - totalPoints.get()) > totalPoints.get() * 0.01) {
            System.out.println("WARNING: Count mismatch > 1%");
          }
        }
      }
    } catch (Exception e) {
      System.err.println("Verification failed: " + e.getMessage());
    }
  }

  private void writeCsv() throws IOException {
    File csvFile = new File(csvOutput);
    csvFile.getParentFile().mkdirs();

    boolean isNewFile = !csvFile.exists();

    try (PrintWriter writer = new PrintWriter(new FileWriter(csvFile, true))) {
      // Write header if new file
      if (isNewFile) {
        writer.println(
            "ts,host,transport,threads,batch,linger_ms,duration_s,"
                + "points_total,points_per_sec,p50_ms,p95_ms,error_rate,"
                + "retries,cpu_pct,heap_mb,gc_count,notes");
      }

      // Calculate metrics
      long duration = (endTime - startTime) / 1000;
      long throughput = duration > 0 ? totalPoints.get() / duration : 0;
      double errorRate =
          totalPoints.get() > 0 ? (double) errorCount.get() / totalPoints.get() * 100 : 0;

      Collections.sort(latencies);
      long p50 = calculatePercentile(latencies, 50);
      long p95 = calculatePercentile(latencies, 95);

      long gcCount = 0;
      for (GarbageCollectorMXBean gc : ManagementFactory.getGarbageCollectorMXBeans()) {
        gcCount += gc.getCollectionCount();
      }

      MemoryMXBean memory = ManagementFactory.getMemoryMXBean();
      long heapMb = memory.getHeapMemoryUsage().getUsed() / (1024 * 1024);

      // Write data row
      SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss");
      String timestamp = sdf.format(new java.util.Date());
      String hostname = java.net.InetAddress.getLocalHost().getHostName();

      writer.printf(
          "%s,%s,%s,%d,%d,%d,%d,%d,%d,%d,%d,%.2f,%d,%.1f,%d,%d,%s%n",
          timestamp,
          hostname,
          transport,
          threads,
          batchSize,
          lingerMs,
          duration,
          totalPoints.get(),
          throughput,
          p50,
          p95,
          errorRate,
          retryCount.get(),
          0.0, // CPU percentage (would need JMX for accurate value)
          heapMb,
          gcCount,
          "" // notes
          );

      System.out.println("Results written to: " + csvFile.getAbsolutePath());
    }
  }

  /** Simple rate limiter for controlling throughput. */
  private static class RateLimiter {
    private final int ratePerSecond;
    private long lastTime = System.nanoTime();
    private double available = 0;

    public RateLimiter(int ratePerSecond) {
      this.ratePerSecond = ratePerSecond;
    }

    public synchronized void acquire(int permits) {
      long now = System.nanoTime();
      double elapsed = (now - lastTime) / 1_000_000_000.0;
      available += elapsed * ratePerSecond;
      available = Math.min(available, ratePerSecond);
      lastTime = now;

      if (available < permits) {
        long sleepMs = (long) ((permits - available) * 1000.0 / ratePerSecond);
        try {
          Thread.sleep(sleepMs);
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
        }
        available = 0;
      } else {
        available -= permits;
      }
    }
  }
}
